import argparse
import time
import pandas as pd
import torch.nn as nn
from matplotlib import pyplot as plt
from torch.nn.utils.parametrizations import weight_norm
from torch.utils.data import DataLoader
from torch.utils.data import Dataset
from tqdm import tqdm
import os
import random
import numpy as np
import torch
os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"


# 设置随机种子以确保结果的可重复性
def set_random_seeds(seed_value):
    random.seed(seed_value)
    np.random.seed(seed_value)
    torch.manual_seed(seed_value)
    if torch.cuda.is_available():
        torch.cuda.manual_seed(seed_value)
        torch.cuda.manual_seed_all(seed_value)  # 为所有GPU设置种子
        # 以下两行设置cuDNN的确定性
        torch.backends.cudnn.deterministic = True
        torch.backends.cudnn.benchmark = False


# 使用示例
set_random_seeds(99)


class StandardScaler():
    # 初始化 StandardScaler 实例，设置 mean 和 std 属性为 0 和 1。这些属性将存储数据的均值和标准差。
    def __init__(self):
        # 设置均值
        self.mean = 0.
        # 设置标准差
        self.std = 1.

    # 计算并存储数据的均值和标准差
    def fit(self, data):
        # 计算均值和标准差，这里的 0 参数表示沿着第一个维度（通常是样本维度）进行计算。
        self.mean = data.mean(0)
        self.std = data.std(0)

    # 对数据进行标准化处理。
    def transform(self, data):
        mean = torch.from_numpy(self.mean).type_as(data).to(data.device) if torch.is_tensor(data) else self.mean
        std = torch.from_numpy(self.std).type_as(data).to(data.device) if torch.is_tensor(data) else self.std
        # 返回标准化后的数据
        return (data - mean) / std

    # 对标准化后的数据进行逆变换，恢复到原始的尺度。
    def inverse_transform(self, data):
        mean = torch.from_numpy(self.mean).type_as(data).to(data.device) if torch.is_tensor(data) else self.mean
        std = torch.from_numpy(self.std).type_as(data).to(data.device) if torch.is_tensor(data) else self.std
        # 检查输入数据 old_data 的最后一个维度（通常是特征维度）是否与 mean 和 std 的形状匹配
        if data.shape[-1] != mean.shape[-1]:
            mean = mean[-1:]
            std = std[-1:]
        # 返回逆变换后的数据，即 (old_data * std) + mean
        return (data * std) + mean


# 可视化损失函数（loss）随迭代次数变化的结果
def plot_loss_data(data):
    # 使用Matplotlib绘制线图
    plt.figure()
    plt.plot(data, label='loss')
    # 添加标题
    plt.title("TCN-LSTM loss results Plot")
    # 显示图例
    plt.legend()
    plt.show()


# Dataset 类是 PyTorch 数据加载框架中的核心组件，用于封装数据集并提供数据的访问接口
class TimeSeriesDataset(Dataset):
    def __init__(self, sequences):
        self.sequences = sequences

    # 返回数据集中序列的数量
    def __len__(self):
        return len(self.sequences)

    # 按索引 index 访问数据集中的单个序列。
    def __getitem__(self, index):
        # 从 sequences 列表中取出索引为 index 的元素，它是一个包含序列和标签的元组，并将它们分别赋值给 sequence 和 label
        sequence, label = self.sequences[index]
        # 将序列和标签转换为 PyTorch 张量（Tensor），并返回。这样，数据就可以被 PyTorch 的数据加载器（DataLoader）直接用于模型训练。
        return torch.Tensor(sequence), torch.Tensor(label)


def create_inout_sequences(input_data, tw, pre_len, config):
    # 创建时间序列数据专用的数据分割器
    inout_seq = []
    L = len(input_data)
    for i in range(L - tw):
        if (i + tw + pre_len) > len(input_data):
            break
        if config.feature == 'MS':
            # train_seq = input_data[i:i + tw]
            train_seq = input_data[:, :-1][i:i + tw]
            train_label = input_data[:, -1:][i + tw:i + tw + pre_len]
        elif config.feature == 'S':
            train_seq = input_data[:, -1:][i:i + tw]
            train_label = input_data[:, -1:][i + tw:i + tw + pre_len]
        else:
            train_seq = input_data[i:i + tw]
            train_label = input_data[i + tw:i + tw + pre_len]
        inout_seq.append((train_seq, train_label))
    return inout_seq


def calculate_mae(y_true, y_pred):
    # 平均绝对误差
    mae = np.mean(np.abs(y_true - y_pred))
    return mae

def calculate_mse(y_true, y_pred):
    mse = np.mean((y_true - y_pred)**2)
    return mse

def create_dataloader(config, device):
    print(">>>>>>>>>>>>>>>>>>>>>>>>>>>>创建数据加载器<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<")
    df = pd.read_csv(config.data_path)
    pre_len = config.pre_len
    train_window = config.window_size

    # 将特征列移到末尾
    target_data = df[[config.target]]
    # 从 df DataFrame 中删除目标特征列。axis=1 表示操作的是列而不是行
    df = df.drop(config.target, axis=1)
    # 使用 pd.concat 函数将 target_data（目标特征列）重新添加到 df DataFrame 的末尾。axis=1 表示沿着列的方向进行连接
    df = pd.concat((df, target_data), axis=1)

    # 获取 df DataFrame 的列名，并存储到 cols_data 中。这里使用 [1:] 切片操作是为了排除第一列
    cols_data = df.columns[1:]
    # 根据 cols_data 中的列名从 df DataFrame 中选择对应的列，存储到 df_data 中。这样 df_data 就包含了除了第一列以外的所有特征列。
    df_data = df[cols_data]

    # 使用 .values 属性将 df_data DataFrame 中的数据转换为一个 NumPy 数组
    true_data = df_data.values

    # 定义标准化优化器
    scaler = StandardScaler()
    # 这行代码使 scaler 根据 true_data 数组中的数据计算均值和标准差
    scaler.fit(true_data)

    # 将数据集分割为训练集、验证集和测试集。
    train_data = true_data[int(0.2 * len(true_data)):]
    # valid_data = true_data[int(0.15 * len(true_data)):int(0.30 * len(true_data))]
    test_data = true_data[:int(0.2 * len(true_data))]
    print("训练集尺寸:", len(train_data), "测试集尺寸:", len(test_data))

    # 进行标准化处理
    train_data_normalized = scaler.transform(train_data)
    test_data_normalized = scaler.transform(test_data)
    # valid_data_normalized = scaler.transform(valid_data)

    train_data_normalized = torch.FloatTensor(train_data_normalized).to(device)
    test_data_normalized = torch.FloatTensor(test_data_normalized).to(device)
    # valid_data_normalized = torch.FloatTensor(valid_data_normalized).to(device)

    train_inout_seq = create_inout_sequences(train_data_normalized, train_window, pre_len, config)
    test_inout_seq = create_inout_sequences(test_data_normalized, train_window, pre_len, config)
    # valid_inout_seq = create_inout_sequences(valid_data_normalized, train_window, pre_len, config)

    # 创建数据集
    train_dataset = TimeSeriesDataset(train_inout_seq)
    test_dataset = TimeSeriesDataset(test_inout_seq)
    # valid_dataset = TimeSeriesDataset(valid_inout_seq)

    train_loader = DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True, drop_last=True)
    test_loader = DataLoader(test_dataset, batch_size=args.batch_size, shuffle=False, drop_last=True)
    # valid_loader = DataLoader(valid_dataset, batch_size=args.batch_size, shuffle=False, drop_last=True)

    print("通过滑动窗口共有训练集数据：", len(train_inout_seq), "转化为批次数据:", len(train_loader))
    print("通过滑动窗口共有测试集数据：", len(test_inout_seq), "转化为批次数据:", len(test_loader))
    # print("通过滑动窗口共有验证集数据：", len(valid_inout_seq), "转化为批次数据:", len(valid_loader))
    print(">>>>>>>>>>>>>>>>>>>>>>>>>>>>创建数据加载器完成<<<<<<<<<<<<<<<<<<<<<<<<<<<")
    # return train_loader, test_loader, valid_loader, scaler
    return train_loader, test_loader, scaler

# nn.Module 的一个子类。
# 在其前向传播（forward）方法中对输入数据 x 进行操作，具体来说，它会移除输入数据最后一个维度上最后 self.chomp_size 个元素
class Chomp1d(nn.Module):
    def __init__(self, chomp_size):
        super(Chomp1d, self).__init__()
        # 指定了要移除的最后一个维度上元素的数量
        self.chomp_size = chomp_size

    def forward(self, x):
        return x[:, :, :-self.chomp_size].contiguous()


# TemporalBlock 是一个包含两个卷积层的序列模型，每个卷积层后面跟着一个 Chomp1d 模块、ReLU 激活函数和 dropout 层
class TemporalBlock(nn.Module):
    # # 初始化 TemporalBlock 实例，接收输入和输出通道数、卷积核大小、步长、扩张率、填充和 dropout 率等参数
    def __init__(self, n_inputs, n_outputs, kernel_size, stride, dilation, padding, dropout=0.2):
        super(TemporalBlock, self).__init__()
        # self.conv1 和 self.conv2 是两个一维卷积层，使用 weight_norm 函数进行权重归一化
        self.conv1 = weight_norm(nn.Conv1d(n_inputs, n_outputs, kernel_size,
                                           stride=stride, padding=padding, dilation=dilation))
        # self.chomp1 和 self.chomp2 是两个 Chomp1d 模块，用于移除卷积层中由于填充而产生的额外元素。
        self.chomp1 = Chomp1d(padding)
        # self.relu1、self.relu2 是激活函数，这里使用的是 ReLU
        self.relu1 = nn.ReLU()
        self.dropout1 = nn.Dropout(dropout)

        self.conv2 = weight_norm(nn.Conv1d(n_outputs, n_outputs, kernel_size,
                                           stride=stride, padding=padding, dilation=dilation))
        self.chomp2 = Chomp1d(padding)
        self.relu2 = nn.ReLU()
        self.dropout2 = nn.Dropout(dropout)

        # self.net 是一个序列模型，包含了上述所有层。
        self.net = nn.Sequential(self.conv1, self.chomp1, self.relu1, self.dropout1,
                                 self.conv2, self.chomp2, self.relu2, self.dropout2)
        # self.downsample 是一个一维卷积层，用于在输入和输出通道数不同时进行下采样
        self.downsample = nn.Conv1d(n_inputs, n_outputs, 1) if n_inputs != n_outputs else None
        self.relu = nn.ReLU()
        self.init_weights()

    # 初始化卷积层的权重。这里使用正态分布初始化权重，标准差为 0.01。如果存在下采样层，也对其进行权重初始化
    def init_weights(self):
        self.conv1.weight.data.normal_(0, 0.01)
        self.conv2.weight.data.normal_(0, 0.01)
        if self.downsample is not None:
            self.downsample.weight.data.normal_(0, 0.01)

    # 定义了模块的前向传播逻辑
    def forward(self, x):
        out = self.net(x)
        # 如果不存在下采样层，则残差连接直接使用原始输入 x；如果存在下采样层，则对输入 x 进行下采样
        res = x if self.downsample is None else self.downsample(x)
        # 将处理后的输出 out 与残差 res 相加，然后通过最终的 ReLU 激活函数，得到最终的输出
        return self.relu(out + res)


# 结合了卷积层和循环层（LSTM）来处理时间序列数据
class TemporalConvNet(nn.Module):
    # 初始化 TemporalConvNet 实例，接收输入通道数、输出维度、预测长度、通道数列表、层数、卷积核大小和 dropout 率等参数
    def __init__(self, num_inputs, outputs, pre_len, num_channels, n_layers, kernel_size=2, dropout=0.2):
        super(TemporalConvNet, self).__init__()
        layers = []
        self.pre_len = pre_len
        self.n_layers = n_layers
        self.hidden_size = num_channels[-2]
        self.hidden = nn.Linear(num_channels[-1], num_channels[-2])
        self.relu = nn.ReLU()
        # self.lstm 是一个 LSTM 层，用于处理序列数据
        self.lstm = nn.LSTM(self.hidden_size, self.hidden_size, n_layers, bias=True,
                            batch_first=True)  # output (batch_size, obs_len, hidden_size)
        # num_levels 是 num_channels 列表的长度，表示模型中卷积层的数量
        num_levels = len(num_channels)
        # 循环构建 TemporalBlock 层
        for i in range(num_levels):
            dilation_size = 2 ** i
            in_channels = num_inputs if i == 0 else num_channels[i - 1]
            out_channels = num_channels[i]
            layers += [TemporalBlock(in_channels, out_channels, kernel_size, stride=1, dilation=dilation_size,
                                     padding=(kernel_size - 1) * dilation_size, dropout=dropout)]

        self.network = nn.Sequential(*layers)
        self.linear = nn.Linear(num_channels[-2], outputs)

    def forward(self, x):
        # 调整输入数据 x 的维度，使其最后一个维度是时间步，以适应卷积层的输入要求
        x = x.permute(0, 2, 1)
        # 将调整后的输入数据通过卷积层网络
        x = self.network(x)
        # 再次调整输出数据的维度，使其第二个维度是时间步，以适应 LSTM 的输入要求
        x = x.permute(0, 2, 1)

        # 获取调整后数据的维度信息。
        batch_size, obs_len, features_size = x.shape  # (batch_size, obs_len, features_size)
        # 通过线性层将卷积层的输出转换为 LSTM 的输入
        xconcat = self.hidden(x)  # (batch_size, obs_len, hidden_size)
        # 初始化一个零张量 H，用于存储 LSTM 的隐藏状态
        H = torch.zeros(batch_size, obs_len - 1, self.hidden_size).to(device)  # (batch_size, obs_len-1, hidden_size)
        # 初始化 LSTM 的初始隐藏状态和细胞状态
        ht = torch.zeros(self.n_layers, batch_size, self.hidden_size).to(
            device)  # (num_layers, batch_size, hidden_size)
        ct = ht.clone()
        # 循环处理每个时间步的数据，通过 LSTM 层更新隐藏状态和细胞状态，并将非最后一个时间步的隐藏状态存储在 H 中
        for t in range(obs_len):
            xt = xconcat[:, t, :].view(batch_size, 1, -1)  # (batch_size, 1, hidden_size)
            out, (ht, ct) = self.lstm(xt, (ht, ct))  # ht size (num_layers, batch_size, hidden_size)
            htt = ht[-1, :, :]  # (batch_size, hidden_size)
            if t != obs_len - 1:
                H[:, t, :] = htt
        H = self.relu(H)  # (batch_size, obs_len-1, hidden_size)
        x = self.linear(H)
        return x[:, -self.pre_len:, :]


# 实现了一个标准的 PyTorch 训练循环，包括前向传播、损失计算、反向传播和参数更新，并在每个 epoch 结束时保存模型
def train(model, args, train_loader):
    start_time = time.time()  # 计算起始时间
    # 将传入的模型赋值给局部变量
    model = model
    # 定义损失函数为均方误差损失
    loss_function = nn.MSELoss()
    # 定义优化器
    optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
    # 从传入的 args 对象中获取训练的轮数（epochs）
    epochs = args.epochs
    model.train()  # 训练模式
    results_loss = []
    # 外层循环遍历每个 epoch。
    for i in tqdm(range(epochs)):
        losss = []
        # 内层循环遍历 train_loader 中的每个批次数据（序列和标签）
        for seq, labels in train_loader:
            # 在每次迭代开始时清零梯度，以避免梯度累积
            optimizer.zero_grad()
            # 执行模型的前向传播，得到预测值
            y_pred = model(seq)
            # 计算预测值和真是标签之间的损失值
            single_loss = loss_function(y_pred, labels)
            # 执行反向传播，计算梯度
            single_loss.backward()
            # 根据计算的梯度更新模型参数
            optimizer.step()
            # 将每个批次的损失值添加到列表中，并从 GPU 转移到 CPU，转换为 NumPy 数组
            losss.append(single_loss.detach().cpu().numpy())
        # 记录和打印损失
        tqdm.write(f"\t Epoch {i + 1} / {epochs}, Loss: {sum(losss) / len(losss)}")
        results_loss.append(sum(losss) / len(losss))

        # 在每个 epoch 结束时保存模型的参数
        torch.save(model.state_dict(), 'save_model.pth')
        # time.sleep(0.1)

    # valid_loss = valid(model, args, scaler, valid_loader)
    # 尚未引入学习率计划后期补上
    # 保存模型

    print(f">>>>>>>>>>>>>>>>>>>>>>模型已保存,用时:{(time.time() - start_time) / 60:.4f} min<<<<<<<<<<<<<<<<<<")
    # 绘制损失数据图表
    plot_loss_data(results_loss)


# 用于在验证集上评估一个已经训练好的 PyTorch 模型的性能。一个标准的验证过程，包括加载模型、设置评估模式、遍历验证集、计算损失和输出平均损失
def valid(model, args, valid_loader):
    lstm_model = model
    # 加载模型进行预测
    lstm_model.load_state_dict(torch.load('save_model.pth'))
    lstm_model.eval()  # 评估模式
    mae_losses = []
    mse_losses = []

    for seq, labels in valid_loader:
        pred = lstm_model(seq)
        mae = calculate_mae(pred.detach().cpu().numpy(), np.array(labels.detach().cpu()))  # MAE误差计算绝对值(预测值  - 真实值)
        mse = calculate_mse(pred.detach().cpu().numpy(), np.array(labels.detach().cpu()))
        mae_losses.append(mae)
        mse_losses.append(mse)

    print("验证集误差MAE:", mae_losses)
    print("验证集误差MSE:", mse_losses)
    mae = sum(mae_losses) / len(mae_losses)
    mse = sum(mse_losses) / len(mse_losses)
    print("验证集平均MAE:", mae)
    print("验证集平均MSE:", mse)


# 用于在测试集上评估一个已经训练好的 PyTorch 模型的性能，并且绘制预测结果与真实值的对比图
def test(model, args, test_loader, scaler):
    # 加载模型进行预测
    mae_losses = []
    mse_losses = []
    model = model
    model.load_state_dict(torch.load('save_model.pth'))
    model.eval()  # 评估模式
    results = []
    labels = []
    for seq, label in test_loader:
        pred = model(seq)
        mae = calculate_mae(pred.detach().cpu().numpy(), np.array(label.detach().cpu()))  # MAE误差计算绝对值(预测值  - 真实值)
        mse = calculate_mse(pred.detach().cpu().numpy(), np.array(label.detach().cpu()))
        mae_losses.append(mae)
        mse_losses.append(mse)
        pred = pred[:, 0, :]
        label = label[:, 0, :]
        pred = scaler.inverse_transform(pred.detach().cpu().numpy())
        label = scaler.inverse_transform(label.detach().cpu().numpy())
        for i in range(len(pred)):
            results.append(pred[i][-1])
            labels.append(label[i][-1])

    print("测试集MAE：", mae_losses)
    print("测试集MSE：", mse_losses)
    average_mae = np.mean(mae_losses)
    average_mse = np.mean(mse_losses)
    print("测试集平均误差MAE:", average_mae)
    print("测试集均方误差MSE:", average_mse)
    # 绘制历史数据
    plt.plot(labels, label='TrueValue')
    # 绘制预测数据
    # 注意这里预测数据的起始x坐标是历史数据的最后一个点的x坐标
    plt.plot(results, label='Prediction')
    # 添加标题和图例
    plt.title("TCN-LSTM test state")
    plt.legend()
    plt.show()

# 检验模型拟合情况，提供了一个直观的方式来检查模型在训练集上的拟合情况，通过绘制预测结果与真实值的对比图
def inspect_model_fit(model, args, train_loader, scaler):
    model = model
    model.load_state_dict(torch.load('save_model.pth'))
    model.eval()  # 评估模式
    results = []
    labels = []

    for seq, label in train_loader:
        pred = model(seq)[:, 0, :]
        label = label[:, 0, :]
        pred = scaler.inverse_transform(pred.detach().cpu().numpy())
        label = scaler.inverse_transform(label.detach().cpu().numpy())
        for i in range(len(pred)):
            results.append(pred[i][-1])
            labels.append(label[i][-1])

    # 绘制历史数据
    plt.plot(labels, label='History')

    plt.plot(results, label='Prediction')

    plt.title("TCN-LSTM inspect model fit state")
    plt.legend()
    plt.show()


def predict(model, args, device):
    df = pd.read_csv(args.data_path)
    if args.feature == 'S':
        # 提取 DataFrame 的最后 args.window_size 行，并转换为 NumPy 数组
        df = df.iloc[:, -1:][-args.window_size:].values  # 转换为nadarry
    elif args.feature == 'MS':
        # 提取 DataFrame 的最后 args.window_size 行，并转换为 NumPy 数组
        df = df.iloc[:, 1:-1][-args.window_size:].values  # 转换为nadarry
    # 定义标准化优化器
    scaler = StandardScaler()
    scaler.fit(df)
    # 使用 scaler 对特征数据进行标准化处理
    pre_data = scaler.transform(df)
    # 将 NumPy 数组转换为 PyTorch 张量，并移动到指定的设备（如 GPU）
    tensor_pred = torch.FloatTensor(pre_data).to(device)
    # 增加一个批次维度，因为 PyTorch 模型通常期望批次维度
    tensor_pred = tensor_pred.unsqueeze(0)  # 单次预测 , 滚动预测功能暂未开发后期补上
    model = model
    model.load_state_dict(torch.load('save_model.pth'))
    model.eval()  # 评估模式
    # 使用模型对处理后的数据进行预测
    pred = model(tensor_pred)[0]
    pred = scaler.inverse_transform(pred.detach().cpu().numpy())
    history_length = len(df[:, -1])
    # 为历史数据生成x轴坐标
    history_x = range(history_length)
    prediction_x = range(history_length - 1, history_length + len(pred[:, -1]) - 1)
    plt.plot(history_x, df[:, -1], label='History')
    plt.plot(prediction_x, pred[:, -1], label='Prediction')
    plt.axvline(history_length - 1, color='red')  # 在图像的x位置处画一条红色竖线
    # 添加标题和图例
    plt.title("TCN-LSTM History and Prediction")
    plt.legend()

if __name__ == '__main__':
    # 创建一个 ArgumentParser 对象，用于处理命令行参数，描述信息设置为 "Time Series forecast"
    parser = argparse.ArgumentParser(description='Time Series forecast')
    parser.add_argument('--model', type=str, default='TCN-LSTM', help="模型持续更新")
    parser.add_argument('--window_size', type=int, default=23, help="时间窗口大小, window_size > pre_len")
    parser.add_argument('--pre_len', type=int, default=12, help="预测未来数据长度")
    # old_data
    parser.add_argument('--shuffle', action='store_true', default=True, help="是否打乱数据加载器中的数据顺序")
    parser.add_argument('--data_path', type=str, default='../old_data/all.csv', help="你的数据数据地址")
    parser.add_argument('--target', type=str, default='RSEI', help='你需要预测的特征列，这个值会最后保存在csv文件里')
    parser.add_argument('--input_size', type=int, default=4, help='你的特征个数不算时间那一列')
    parser.add_argument('--feature', type=str, default='MS', help='[M, S, MS],多元预测多元,单元预测单元,多元预测单元')
    parser.add_argument('--model_dim', type=list, default=[64, 128, 256], help='这个地方是这个TCN卷积的关键部分,它代表了TCN的层数我这里输'
                                                                              '入list中包含三个元素那么我的TCN就是三层，这个根据你的数据复杂度来设置'
                                                                              '层数越多对应数据越复杂但是不要超过5层')

    # learning
    parser.add_argument('--lr', type=float, default=0.002, help="学习率")
    parser.add_argument('--drop_out', type=float, default=0.2, help="随机丢弃概率,防止过拟合")
    parser.add_argument('--epochs', type=int, default=80, help="训练轮次")
    parser.add_argument('--batch_size', type=int, default=4, help="批次大小")
    parser.add_argument('--save_path', type=str, default='models')

    # model
    parser.add_argument('--hidden_size', type=int, default=64, help="隐藏层单元数")
    parser.add_argument('--kernel_sizes', type=int, default=3)
    parser.add_argument('--laryer_num', type=int, default=1)
    # device
    parser.add_argument('--use_gpu', type=bool, default=True)
    parser.add_argument('--device', type=int, default=0, help="只设置最多支持单个gpu训练")

    # option
    parser.add_argument('--train', type=bool, default=True)
    parser.add_argument('--valid', type=bool, default=True)
    parser.add_argument('--test', type=bool, default=True)
    parser.add_argument('--predict', type=bool, default=True)
    parser.add_argument('--inspect_fit', type=bool, default=True)
    parser.add_argument('--lr-scheduler', type=bool, default=True)
    # 解析命令行输入的参数，并将它们存储在 args 对象中
    args = parser.parse_args()

    if isinstance(args.device, int) and args.use_gpu:
        device = torch.device("cuda:" + f'{args.device}')
    else:
        device = torch.device("cpu")
    print("使用设备:", device)
    # 调用 create_dataloader 函数，根据 args 和 device 创建训练、测试和验证的数据加载器 (DataLoader) 以及数据标准化器 (scaler)
    # train_loader, test_loader, valid_loader, scaler = create_dataloader(args, device)
    train_loader, test_loader, scaler = create_dataloader(args, device)

    if args.feature == 'MS' or args.feature == 'S':
        args.output_size = 1
    else:
        args.output_size = args.input_size

    # 实例化模型
    try:
        print(f">>>>>>>>>>>>>>>>>>>>>>>>>开始初始化{args.model}模型<<<<<<<<<<<<<<<<<<<<<<<<<<<")
        model = TemporalConvNet(args.input_size, args.output_size, args.pre_len, args.model_dim, args.laryer_num,
                                args.kernel_sizes).to(device)
        print(f">>>>>>>>>>>>>>>>>>>>>>>>>开始初始化{args.model}模型成功<<<<<<<<<<<<<<<<<<<<<<<<<<<")
    except:
        print(f">>>>>>>>>>>>>>>>>>>>>>>>>开始初始化{args.model}模型失败<<<<<<<<<<<<<<<<<<<<<<<<<<<")

    # 训练模型
    if args.train:
        print(f">>>>>>>>>>>>>>>>>>>>>>>>>开始{args.model}模型训练<<<<<<<<<<<<<<<<<<<<<<<<<<<")
        train(model, args, train_loader)
    # if args.valid:
    #     print(f">>>>>>>>>>>>>>>>>>>>>>>>>开始{args.model}模型训练<<<<<<<<<<<<<<<<<<<<<<<<<<<")
    #     valid(model, args, valid_loader)
    if args.test:
        print(f">>>>>>>>>>>>>>>>>>>>>>>>>开始{args.model}模型测试<<<<<<<<<<<<<<<<<<<<<<<<<<<")
        test(model, args, test_loader, scaler)
    if args.inspect_fit:
        print(f">>>>>>>>>>>>>>>>>>>>>>>>>开始检验{args.model}模型拟合情况<<<<<<<<<<<<<<<<<<<<<<<<<<<")
        inspect_model_fit(model, args, train_loader, scaler)
    if args.predict:
        print(f">>>>>>>>>>>>>>>>>>>>>>>>>预测未来{args.pre_len}条数据<<<<<<<<<<<<<<<<<<<<<<<<<<<")
        predict(model, args, device)
    plt.show()